import anndata as ad
import numpy as np
import pandas as pd
from typing import Optional, Union, List, Dict
import scanpy as sc

class CustomAnnData(ad.AnnData):
    def __init__(self, *args, **kwargs):
        # 调用父类的初始化方法
        super().__init__(*args, **kwargs)
        
    def normalize_and_log(self, 
                         target_sum: Optional[float] = 1e4,
                         min_genes: int = 200,
                         min_cells: int = 3):
        """
        标准化和对数转换的组合方法
        """
        # 过滤细胞和基因
        sc.pp.filter_cells(self, min_genes=min_genes)
        sc.pp.filter_genes(self, min_cells=min_cells)
        
        # 标准化
        sc.pp.normalize_total(self, target_sum=target_sum)
        
        # 对数转换
        sc.pp.log1p(self)
        
        # 记录处理状态
        self.uns['normalized'] = True
        return self
    
    def add_metadata(self, 
                    metadata: pd.DataFrame,
                    join: str = 'left'):
        """
        添加细胞元数据
        """
        if not isinstance(metadata, pd.DataFrame):
            raise TypeError("metadata must be a pandas DataFrame")
            
        # 合并元数据
        self.obs = self.obs.join(metadata, how=join)
        return self
    
    def quick_cluster(self,
                     n_pcs: int = 30,
                     n_neighbors: int = 15,
                     resolution: float = 0.5,
                     random_state: int = 42):
        """
        快速聚类分析流程
        """
        # 计算高变基因
        sc.pp.highly_variable_genes(self)
        
        # PCA
        sc.tl.pca(self, n_comps=n_pcs, random_state=random_state)
        
        # 计算邻居图
        sc.pp.neighbors(self, n_neighbors=n_neighbors, random_state=random_state)
        
        # Leiden聚类
        sc.tl.leiden(self, resolution=resolution, random_state=random_state)
        
        # UMAP可视化
        sc.tl.umap(self, random_state=random_state)
        
        return self
    
    def find_markers(self,
                    groupby: str = 'leiden',
                    n_genes: int = 25,
                    method: str = 'wilcoxon') -> Dict[str, pd.DataFrame]:
        """
        找到每个群的marker基因
        """
        sc.tl.rank_genes_groups(self, 
                              groupby=groupby,
                              method=method)
        # 获取结果
        markers = {}
        groups = self.obs[groupby].unique()
        
        for group in groups:
            markers[group] = sc.get.rank_genes_groups_df(self,
                                                       group=group,
                                                       n_genes=n_genes)
        return markers
    
    @property
    def is_normalized(self) -> bool:
        """
        检查数据是否已经标准化
        """
        return self.uns.get('normalized', False)
    
    def subset_cells(self, 
                    cell_indices: Union[List[int], np.ndarray]) -> 'CustomAnnData':
        """
        创建细胞子集
        """
        return CustomAnnData(self[cell_indices].copy())

# 使用示例
if __name__ == "__main__":
    # 创建示例数据
    data = np.random.rand(1000, 100)  # 1000个细胞，100个基因
    custom_adata = CustomAnnData(data)
    
    # 添加一些基本的观察信息
    custom_adata.obs['condition'] = np.random.choice(['A', 'B'], size=1000)
    
    # 使用自定义方法
    custom_adata.normalize_and_log()
    
    # 添加额外的元数据
    metadata = pd.DataFrame({
        'batch': np.random.choice(['1', '2', '3'], size=1000)
    }, index=custom_adata.obs.index)
    custom_adata.add_metadata(metadata)
    
    # 运行聚类分析
    custom_adata.quick_cluster()
    
    # 找到marker基因
    markers = custom_adata.find_markers()
    
    # 创建子集
    subset = custom_adata.subset_cells(np.arange(100))  # 前100个细胞
